#!/usr/bin/env python
import csv, re, random
import numpy as np
import pyphon
import pyphon.bitERCs as be

def stressWeight(output):
    return output.replace('[[L', '1L').replace('[[H', '1H')\
            .replace('L]]', '1L').replace('H]]', '1H').replace('L', '2L')\
            .replace('H', '2H').replace('l', '0L').replace('h', '0H')\
            .replace('11', '1').replace('12', '1').replace('[', '')\
            .replace(']', '')

def stress(output):
    return stressWeight(output).replace('H', '').replace('L', '')

def inputOf(output):
    inputSigma = ('#', 'l', 'h')
    return ''.join([ch for ch in output.lower() if ch in inputSigma])

def typDistinctOutputsStressWeight(typfile):
    osets = set()
    r = csv.reader(open(typfile, 'r'))
    
    nlangs = 0
    for row in r:
        if row and row[0] == 'Surface forms:':
            nlangs += 1
            forms = frozenset([stressWeight(output) for output in\
                row[1][1:-1].split(', ')])
            osets.add(forms)

    print '%d strong languages, %d surface languages' % (nlangs, len(osets))
    
    return osets

def typDistinctOutputsStress(typfile):
    osets = set()
    r = csv.reader(open(typfile, 'r'))
    
    for row in r:
        if row and row[0] == 'Surface forms:':
            forms = frozenset([stress(output) for output in\
                row[1][1:-1].split(', ')])
            osets.add(forms)
    
    return osets

def surfaceTypology(osets):
    surfTyp = {}
    lnum = 1
    for oset in osets:
        lname = 'S%04d' % lnum
        lnum += 1
        slang = {}
        for output in oset:
            slang[inputOf(output)] = output
        surfTyp[lname] = slang

        # for tessmo inputs
        #assert(len(slang) == 62)

    return surfTyp

def ambiguityStats(tabfile, statsfile):
    """
    Reads tableaux from tabfile and write ambiguity statistics to statsfile.
    Both are file names.
    """
    r = csv.reader(open(tabfile, 'r'))
    tableaux = []

    lastInput = None
    curTableau = None
    for row in r:
        if row and row[0]:
            input, output = row[0], row[1]
            if ' | ' in output:
                # shouldn't happen
                raise ValueError('Found |')
            if input != lastInput:
                if curTableau:
                    tableaux.append(curTableau)
                curTableau = {'input': input, 'candidates':[]}
            lastInput = input
            curTableau['candidates'].append(output)

    w = csv.writer(open(statsfile, 'w'))
    w.writerow(('input', 'inputLen', 'candidates', 'distinctStressWeight',
        'distinctStress'))
    for tab in tableaux:
        input = tab['input']
        stressWeightOutputs = [stressWeight(o) for o in tab['candidates']]
        stressOutputs = [stress(o) for o in tab['candidates']]
        w.writerow((input, len(input)-2, len(tab['candidates']),
            len(set(stressWeightOutputs)), len(set(stressOutputs))))

def surfaceTableaux_ebits(tabfile):
    r = csv.reader(open(tabfile, 'r'))
    tableaux = {}

    lastInput = None
    curTableau = None
    for row in r:
        if row and row[0]:
            input, output, violns = row[0], row[1], [int(v) for v in row[3:]]
            if ' | ' in output:
                # shouldn't happen
                raise ValueError('Found |')
            if input != lastInput:
                if curTableau:
                    tableaux[lastInput] = curTableau
                curTableau = {'candidates':{}, 'input':input}
            lastInput = input
            curTableau['candidates'][output] = tuple(violns)
            curTableau['k'] = len(violns)
    tableaux[lastInput] = curTableau

    for input, tab in tableaux.iteritems():
        surfaceCandidates = {}
        tab['ercs'] = {}
        for o,v in tab['candidates'].iteritems():
            swo = stressWeight(o)
            if swo not in surfaceCandidates:
                surfaceCandidates[swo] = []
            surfaceCandidates[swo].append(o)

            ercs = []
            for loser, loserV in tab['candidates'].iteritems():
                if v == loserV:
                    continue
                ercs.append(be.ercBit(v, loserV))
            tab['ercs'][o] = frozenset(ercs)

        tab['surfaceCandidates'] = surfaceCandidates

    for input, tab in tableaux.iteritems():
        del tab['candidates']

    return tableaux

def surfaceTableaux_pyphon(tabfile):
    r = csv.reader(open(tabfile, 'r'))
    tableaux = {}

    lastInput = None
    curTableau = None
    firstRow = r.next()
    con = firstRow[3:]
    print 'CON:', con
    for row in r:
        if row and row[0]:
            input, output, violns = row[0], row[1], [int(v) for v in row[3:]]
            if ' | ' in output:
                # shouldn't happen
                raise ValueError('Found |')
            if input != lastInput:
                if curTableau:
                    tableaux[lastInput] = curTableau
                curTableau = {'candidates':{}, 'input':input}
            lastInput = input
            curTableau['candidates'][output] = tuple(violns)
            curTableau['k'] = len(violns)
    tableaux[lastInput] = curTableau

    for input, tab in tableaux.iteritems():
        surfaceCandidates = {}
        tab['ercs'] = {}
        for o,v in tab['candidates'].iteritems():
            swo = stressWeight(o)
            if swo not in surfaceCandidates:
                surfaceCandidates[swo] = []
            surfaceCandidates[swo].append(o)

            #ercs = []
            #for loser, loserV in tab['candidates'].iteritems():
            #    if v == loserV:
            #        continue
            #    ercs.append(be.ercBit(v, loserV))
            try:
                tab['ercs'][o] = pyphon.pycbg.impliedERCSet(zip(con, v), [zip(con,
                    loserV) for loser, loserV in tab['candidates'].iteritems() \
                        if v != loserV])
            except pyphon.pycbg.InconsistentERCSetError, err:
                print 'Tableau %s is b0rked! Putative winner ain\'t: %s' %\
                    (input, o)
                print 'ERC set:'
                print err.eSet.stringERCs()
                print 'Candidates:'
                return tab['candidates']

        tab['surfaceCandidates'] = surfaceCandidates

    #for input, tab in tableaux.iteritems():
    #    del tab['candidates']

    return tableaux

def learnFromSurfaceLang(surfTyp, surfTabs, surfLang, learnFn, report=True):
    print 'Learning from surface language %s\n' % surfLang
    slang = surfTyp[surfLang]

    hypothesis =  learnFn(slang, surfTabs)
    print '\nFinished learning from language %s.' % surfLang
    if report:
        return reportOnHypothesis(hypothesis, slang, surfTabs)

def reportOnHypothesis(hypothesis, slang, surfTabs):
    print 'Report:'
    nTabsWithAmbig = 0
    noHulls = 0
    ambigProduct = 1
    for input, tab in surfTabs.iteritems():
        k = tab['k']
        print '\tFor input %s' % input
        print '\tSurface output:', slang[input]
        print '\tOriginally ambiguous for %d parses:' %\
            len(tab['surfaceCandidates'][slang[input]])

        for ambigOut in tab['surfaceCandidates'][slang[input]]:
            print '\t\t%s' % ambigOut

        consistentOuts = []
        for o, es in tab['ercs'].iteritems():
            if be.eBitsConsistent(es | hypothesis, k):
                consistentOuts.append(o)

        ambigProduct *= len(consistentOuts)
        print '\t%d parse%s %s consistent with hypothesis:' %\
                (len(consistentOuts), 's' if len(consistentOuts) > 1 else '',
                 'are' if len(consistentOuts) > 1 else 'is')
        for co in consistentOuts:
            print '\t\t%s' % co

        if len(consistentOuts) > 1:
            disj = be.eBitDisj([tab['ercs'][ao] for ao in
                consistentOuts], k)
            if not disj.issubset(hypothesis):
                print '\t*** But their hull is not in the hypothesis!'
                noHulls += 1

        print
        if len(consistentOuts) > 1:
            nTabsWithAmbig += 1

    print 'After learning, there are still %d tableaux with ambiguous outputs.'\
            % nTabsWithAmbig
    print 'The language remains %d-ways ambiguous (log = %f).' % (ambigProduct,
            np.log2(ambigProduct))
    print 'There were %d tableaux whose hypothesis-viable ambiguous outputs\' ' \
          'disjunctions where not contained in the hypothesis.' % noHulls

    return nTabsWithAmbig, ambigProduct


def ignorantLearner(slang, surfTabs):
    candsToDo = []
    nextCandsToDo = []
    k = 0
    for input in slang:
        tab = surfTabs[input]
        k = tab['k']
        surfOut = slang[input]
        ambig = tab['surfaceCandidates'][surfOut]
        ambigERCs = [tab['ercs'][ao] for ao in ambig]
        candsToDo.append(ambigERCs)

    candsToDoLens = sorted([len(x) for x in candsToDo])
    random.shuffle(candsToDo)
    hypothesis = set()

    while 1:
        print 'len(hypothesis) =', len(hypothesis)
        print len(candsToDo), 'candidates to do'

        if not candsToDo:
            print 'Lets go again'
            nextCandsToDoLens = sorted([len(x) for x in nextCandsToDo])
            if nextCandsToDoLens == candsToDoLens:
                print 'Finished in a most delightful way!'
                break
            if not nextCandsToDo:
                print 'Finished in a less delightful way!'
                break

            candsToDo = list(nextCandsToDo)
            candsToDoLens = sorted([len(x) for x in candsToDo])
            del nextCandsToDo[:]

        cand = candsToDo.pop()

        viableCand = []
        for aoERCs in cand:
            if be.eBitsConsistent(aoERCs | hypothesis, k):
                viableCand.append(aoERCs)

        if len(viableCand) > 1:
            nextCandsToDo.append(viableCand)
        elif len(viableCand) == 1:
            hypothesis |= viableCand[0]
        else:
            print 'Bad candidate:'
            print cand
            print viableCand
            raise 'OH SHIT'

        #if not disj.issubset(hypothesis):
        #print 'adding disj:', disj
        #print 'There be inputs to do:', inputsToDo

    return hypothesis


def disjunctionLearner(slang, surfTabs):
    candsToDo = []
    nextCandsToDo = []
    k = 0
    for input in slang:
        tab = surfTabs[input]
        k = tab['k']
        surfOut = slang[input]
        ambig = tab['surfaceCandidates'][surfOut]
        ambigERCs = [tab['ercs'][ao] for ao in ambig]
        candsToDo.append(ambigERCs)

    candsToDoLens = sorted([len(x) for x in candsToDo])
    random.shuffle(candsToDo)
    hypothesis = set()

    while 1:
        print 'len(hypothesis) =', len(hypothesis)
        print len(candsToDo), 'candidates to do'
        #print 'candsToDoLens:', candsToDoLens

        if not candsToDo:
            print 'Lets go again'
            nextCandsToDoLens = sorted([len(x) for x in nextCandsToDo])
            if nextCandsToDoLens == candsToDoLens:
                print 'Finished in a most delightful way!'
                break
            if not nextCandsToDo:
                print 'Finished in a less delightful way!'
                break
            candsToDo = list(nextCandsToDo)
            candsToDoLens = sorted([len(x) for x in candsToDo])
            del nextCandsToDo[:]

        cand = candsToDo.pop()
        viableCand = []
        for aoERCs in cand:
            if be.eBitsConsistent(aoERCs | hypothesis, k):
                viableCand.append(aoERCs)

        if len(viableCand) > 1:
            disj = be.eBitDisj(viableCand, k)
            nextCandsToDo.append(viableCand)
            #print 'Added viableCand to nextCandsToDo.'
            #print '\tlen(viableCand) =', len(viableCand)
            #print '\tlen(nextCandsToDo) =', len(nextCandsToDo)
        elif len(viableCand) == 1:
            disj = viableCand[0]
        else:
            # should ne'er ha'en
            raise ValueError('No hypotheses left! WTF, Mate?')

        #if not disj.issubset(hypothesis):
        #print 'adding disj:', disj
        hypothesis |= disj
        #print 'There be inputs to do:', inputsToDo

    return hypothesis


def learnerStats(statfile, styp, stabs):
    f = open(statfile, 'w')
    w = csv.writer(f)
    w.writerow(('learner', 'surfLang', 'nTabsAmbig', 'ambigProduct'))
    for slang in styp.iterkeys():
        tI, pI = learnFromSurfaceLang(styp, stabs, slang, ignorantLearner)
        w.writerow(('ignorant', slang, tI, pI))

        tD, pD = learnFromSurfaceLang(styp, stabs, slang, disjunctionLearner)
        w.writerow(('disjunction', slang, tD, pD))

        f.flush()

    f.close()
    

def typAmbiguityStats(stp, stabs):
    pass


def allInputsOfLen(minLen, maxLen):
    for il in xrange(minLen, maxLen+1):
        for n in xrange(2**il):
            yield '\\#%s\\#' % np.binary_repr(n).zfill(il).replace('0',
                    'l').replace('1', 'h')
